# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: leon

import logging
import time
import threading
from concurrent.futures import ThreadPoolExecutor

import tornado.gen
import requests
from google.protobuf import json_format

from coin.exchange.uniswap_v3.kr_rest.native_private_client_v2 import UniswapV3NativePrivateClientV2
from coin.proto.coin_query_pb2 import (
    AccountBalance,
    CurrencyBalance,
)

import coin.exchange.base.kr_rest.private_client_base as privcb
from coin.exchange.uniswap_v3.kr_rest.product import UniswapProduct
from coin.exchange.base.kr_rest.rest_client_base import RestUpdate
from coin.exchange.uniswap_v2.kr_rest.private_client_v2 import UniswapPrivateParser


class UniswapPrivateClientLocal(privcb.PrivateClientBase):
    ProductType = UniswapProduct

    def __init__(
        self,
        wallet_address,
        og_info,
        blockchain_name="ethereum",
        mea="Spot.Uniswap.v3",
    ):
        privcb.PrivateClientBase.__init__(self)
        self._wallet_address = wallet_address
        self._blockchain_name = blockchain_name
        self._mea = mea
        self._currency_list = [curr.native_currency for curr in og_info.currencies]
        if 'ETH' not in self._currency_list:
            self._currency_list.append('ETH')
        if 'WETH' not in self._currency_list:
            self._currency_list.append('WETH')
        self._symbol_list = [prod.native_symbol for prod in og_info.products]
        self._web3_api_client = UniswapV3NativePrivateClientV2(blockchain_name=blockchain_name)
        self._executor = ThreadPoolExecutor(max_workers=2)
        self._fire_lock_map = {symbol: threading.Lock() for symbol in self._symbol_list}

    @staticmethod
    def parse_balance_from_dict(account_balances_dict):
        account_balances = account_balances_dict
        currency_balance_list = [
            CurrencyBalance(
                currency=currency,
                currency_native=currency,
                available=balance["available"],
                hold=balance["hold"],
                total=balance["hold"] + balance["available"],
            ) for currency, balance in account_balances.items()
        ]
        return AccountBalance(exchange="Uniswap", market_type="Spot", each_balance=currency_balance_list)

    def _async_account_balance_impl(self):
        if self._executor:
            fut = self._executor.submit(self._web3_api_client.get_balances, self._wallet_address, self._currency_list)
            return fut
        return None

    @tornado.gen.coroutine
    def async_query_account_balance_impl(self):
        req_timestamp = time.time()
        fut = self._async_account_balance_impl()
        account_balance_dict = yield fut

        resp_timestamp = time.time()
        update = RestUpdate(
            req_timestamp,
            resp_timestamp,
            account_balance_dict,
            requests.codes.ok,
        )
        update.msg = self.parse_balance_from_dict(account_balance_dict)
        return update

    def _async_query_nft_token_info(self, symbol: str):
        if self._executor:
            fut = self._executor.submit(self._web3_api_client.get_nft_token_info, self._wallet_address, symbol)
            return fut
        return None

    @tornado.gen.coroutine
    def async_query_nft_token_info(self):
        nft_token_infos = {}
        for symbol in self._symbol_list:
            fut = self._async_query_nft_token_info(symbol)
            tokenid_to_nft_info = yield fut
            nft_token_infos[symbol] = {"tokenid_to_nft_info": tokenid_to_nft_info}
        return nft_token_infos

    def _withdraw_liquidity(self, symbol, execution_lock):
        try:
            return_status, return_tx_hash= self._web3_api_client.withdraw_liquidity(
                wallet_address=self._wallet_address,
                symbol=symbol,
            )            
        except Exception:
            logging.exception('Exception raised from _web3_api_client.query_withdraw_liquidity.')
            if execution_lock and execution_lock.locked():
                execution_lock.release()
            return None
        if execution_lock:
            assert execution_lock.locked()
            execution_lock.release()
        return {"status": return_status, "tx_hash": return_tx_hash}

    def _provide_liquidity(
        self,
        symbol,
        execution_lock,
        price_upper,
        price_lower,
        **kwargs
    ):
        try:
            return_status, return_tx_hash, fill_price = self._web3_api_client.provide_liquidity(
                wallet_address=self._wallet_address,
                symbol=symbol,
                min_price=price_lower,
                max_price=price_upper,
            )
        except Exception:
            logging.exception('Exception raised from _web3_api_client.query_provide_liquidity.')
            if execution_lock and execution_lock.locked():
                execution_lock.release()
            return None
        if execution_lock:
            assert execution_lock.locked()
            execution_lock.release()
        return {"status": return_status, "tx_hash": return_tx_hash, "fill_price": fill_price}

    def query_liquidity_pool_submit_impl(self, product, action_type, execution_lock, **kwargs):
        if self._executor:
            if action_type == "PROVIDE":
                return self._executor.submit(self._provide_liquidity, product.native_symbol, execution_lock, **kwargs)
            elif action_type == "WITHDRAW":
                return self._executor.submit(self._withdraw_liquidity, product.native_symbol, execution_lock)
            else:
                raise ValueError(product, action_type)
        return None

    def _query_submit_impl(self, symbol, side, price, qty):
        lock = self._fire_lock_map[symbol]
        if lock.acquire(False):
            try:
                return_status, return_tx_hash, fill_price = self._web3_api_client.submit_order(
                    wallet_address=self._wallet_address,
                    symbol=symbol,
                    price=float(price),
                    qty=float(qty),
                    side=side.name,
                )
            except Exception:
                logging.exception('Exception raised from _web3_api_client.query_submit_order.')
                lock.release()
                return None
            lock.release()
            return {"status": return_status, "tx_hash": return_tx_hash, "fill_price": fill_price}
        else:
            logging.warning(f"order rejected: {symbol}")
            return None

    def query_submit_impl(self, product, side, price, qty):
        symbol = product.native_symbol
        if self._executor:
            return self._executor.submit(self._query_submit_impl, symbol, side, price, qty)
        else:
            return None

    def query_cancel_product_impl(self, product):
        raise NotImplementedError()

    def query_nft_token_info(self, product):
        raise NotImplementedError()

    def query_fills_impl(self, product):
        raise NotImplementedError()

    def query_list_orders_impl(self, product):
        raise NotImplementedError()

    def query_account_balance_impl(self):
        raise NotImplementedError()
